import argparse, sys, signal, time
import socket, select
import Queue
import sqlite3 as db

POLL_TIMEOUT = 400
RECV_BUFFSIZE = 1024 * 16

def enum(**enums):
  return type('Enum', (), enums)

def handleArgs():
  parser = argparse.ArgumentParser(description='Network Layer Agent')
  parser.add_argument('--ip', action='store', dest='ip', default='0.0.0.0', help="IP address for this agent. By default we bind to all.")
  parser.add_argument('--db', action='store', dest='dbname', default='flows.db', help="Filename for flow database.")
  args = parser.parse_args()
  return args

def cleanup(server_sockets, portlist, dbconn, dbcurs):
  for p in portlist:
    print >> sys.stderr, 'Closing socket bound to port %d' %  p
    server_sockets[p].close()
  # We don't really need to commit, but we will put it here anyway.
  dbconn.commit()  
  dbcurs.close()
  sys.exit()
  
State = enum(RECV=1, WAIT=2, SEND=3)

class Client:
  def __init__(self):
    self.state = State.RECV
    self.bytesrecv = 0
    self.bytessent = 0
    self.sendtimer = 0

if __name__ == '__main__':
  args = handleArgs()
  print 'Responding as %s' % args.ip

  # Zero out some memory we can copy to send
  empty_bytes = bytearray(1000000)
  
  # Connect to database.
  dbconn = db.connect(args.dbname)
  dbcurs = dbconn.cursor()
  
  # Determine a list of what ports to listen to:
  # Eliminate duplicates
  portlist = list(set([6789,6542]))
  portlist.sort()

  # Set up/bind to the sockets we will be listening to.
  server_sockets = {}
  outputs = []
  inputs = []
  for p in portlist:
    server_sockets[p] = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sockets[p].setblocking(0)
    server_address = (args.ip, p)
    print >> sys.stderr, 'Binding to %s port %s' % server_address
    server_sockets[p].bind(server_address)
    server_sockets[p].listen(5)
    inputs.append(server_sockets[p])

  
  # Outgoing message queues
  message_queues = {}
  # Keep track of how many bytes we've received and sent to each client
  bytesrecvd = {}
  bytessent  = {}
  # Keep track of when we should send
  sendtimes = {}
  # Keep track of what state we're in
  connstate = {}
  
  # Keep track of connections
  clients = {}
 
  # Till the end of time
  while inputs:
    # Check time once per loop? 
    loop_time = time.time()
    print >> sys.stderr, 'TIME: ', loop_time
    for c in clients:
      if clients[c].state == State.WAIT and loop_time >= clients[c].sendtimer:
        print >> sys.stderr, '\nDone waiting, ready to send' 
        clients[c].state = State.SEND
    
    try:
      print >> sys.stderr, '\nWaiting for next connection'
      readable, writable, exceptional = select.select(inputs, outputs, inputs)

      for socket in readable:
        if socket in server_sockets.values():
          # We got a new connection
          connection, client_address = socket.accept()
          our_address = connection.getsockname()
          print >> sys.stderr, 'New connection from:%s->%s' % (client_address, our_address)
          connection.setblocking(0)
          inputs.append(connection)
          # Allocate a queue for data we want to send to client
          message_queues[connection] = Queue.Queue()
          
          # Keep track of state for client
          clients[connection] = Client()
          
        else:
          # Client has sent us data
          data = ''
          data = socket.recv(RECV_BUFFSIZE)
          if data:
            print >> sys.stderr, 'Received %d bytes from %s' % (len(data), socket.getpeername())
            clients[socket].bytesrecv = clients[socket].bytesrecv + len(data)
            if clients[socket].bytesrecv >= 16:
              print >> sys.stderr, 'Starting timer and going into wait time'
              clients[socket].bytesrecv = 0
              clients[socket].sendtimer = loop_time + 5
              clients[socket].state = State.WAIT
          else:
            # Client must have disconnected...
            print >> sys.stderr, 'Closing connection with', client_address, 'after reading no data'
            # Clean up:
            if socket in outputs:
              outputs.remove(socket)
            inputs.remove(socket)
            socket.close()
            del message_queues[socket]
    
      # Send out messages
      for socket in writable:
        # Check the delay timer to see if it's time to send.
        #try:
        #  next_mesg = message_queues[socket].get_nowait()      
        #except Queue.Empty:
        #  print >> sys.stderr, 'Nothing to send to', socket.getpeername()
        #  outputs.remove(socket)
        #else:
        #  socket.send(next_mesg)
        
        # Check if we've received all the bytes we're suppose to, and then see if the timer's expired.
        if clients[socket].state == State.SEND:
          sent = socket.send('b'*80)
          print >> sys.stderr, 'Sent ', sent, ' bytes'
          clients[socket].bytessent = sent + client[socket].bytessent
          if clients[socket].bytessent >= 16:
            # We've completed a recv, wait, send round. 
            bytessent[socket] = 0
            connstate[socket] = State.RECV

      # Handle errors
      for socket in exceptional:
        print >> sys.stderr, 'Handing exceptional case for', socket.getpeername()
        inputs.remove(socket)
        if socket in outputs:
          outputs.remove(socket)
        socket.close()
        del message_queues[socket]          

    except KeyboardInterrupt:
      # Break down
      print >> sys.stderr, 'Cleaning up as best we can'
      cleanup(server_sockets, portlist, dbconn, dbcurs)
  
  # We're out of the loop...?
  print "We shouldn't get here."
  cleanup(server_sockets, portlist, dbconn, dbcurs)
